/*!
    Copright: Kai "DrHalan" Mast - www.freakybytes.org

    This file is part of OpenNetwork. See OpenNetwork.h for license information.
*/

#include "CUDPPort.h"

#include "include/OpenNetwork/OpenNetwork.h"

#include <errno.h>

namespace onet
{

    CUDPPort::CUDPPort(INetworkCallback* callback_, u16 maxConnections_, u16 port_, bool listen_, bool IPv6_)
            : INetworkPort(callback_, maxConnections_, port_, listen_, IPv6_)
    {
        if (maxPeers)
            peers.init(maxPeers);
    }

    CUDPPort::~CUDPPort()
    {
        this->disconnectAll();

        //!Important to close the file descriptor afterwards
        close(socketfd);
    }

    bool CUDPPort::init()
    {
        //! The actual network code
        //! We've to check wheter the port we use uses IPv6 or not and differentiate
        if (this->isIPv6())
            socketfd = socket(PF_INET6, SOCK_DGRAM, 0);
        else
            socketfd = socket(PF_INET, SOCK_DGRAM, 0);

        if (socketfd == -1)
        {
            perror("socket");
            return false;
        }

        const s32 y = 1;

#ifndef WINDOWS
        setsockopt( socketfd, SOL_SOCKET, SO_REUSEADDR, &y, sizeof(s32) );
#else
        setsockopt( socketfd, SOL_SOCKET, SO_REUSEADDR, (char*)&y, sizeof(s32) );
#endif

        //! If we're not listening we don't have to bind either
        if (this->isIPv6())
        {
            struct sockaddr_in6 myaddr;

            myaddr.sin6_family = AF_INET6;

            myaddr.sin6_port = htons( portNumber );
            myaddr.sin6_addr = in6addr_any;

            if (bind(socketfd, (struct sockaddr*)&myaddr, sizeof(myaddr)) == -1)
            {
                perror("bind");
                return false;
            }
        }
        else
        {
            struct sockaddr_in myaddr;

            myaddr.sin_family = AF_INET;

            myaddr.sin_port = htons( portNumber );
            myaddr.sin_addr.s_addr = INADDR_ANY;

            if (bind(socketfd, (struct sockaddr*)&myaddr, sizeof(myaddr)) == -1)
            {
                perror("bind");
                return false;
            }
        }

        //! Make the socket non blocking
#ifndef WINDOWS
        s32 save_fd = fcntl( socketfd, F_GETFL );

        if (save_fd == -1)
        {
            perror("fcntl");
            return false;
        }
        else
        {
            save_fd |= O_NONBLOCK;
            fcntl( socketfd , F_SETFL, save_fd );
        }
#else
        u_long iMode = 1;

        if (ioctlsocket(socketfd, FIONBIO, &iMode) != 0)
        {
            perror("ioctlsocket");
            return false;
        }
#endif

        this->getCallback()->onCreate(this);

        return true;
    }

    SNetworkAddress CUDPPort::getAddress(const PeerID id)
    {
        if (!peers[id])
        {
            if (this->isIPv6())
                return SNetworkAddress("::/128", 0, true);
            else
                return SNetworkAddress("0.0.0.0", 0, false);
        }

        return peers[id]->address;
    }

//! Connect to an address
    //! Returns false if already connected
    //! Address will be deleted itself
    bool CUDPPort::connect(const SNetworkAddress& address)
    {
        if (this->isConnected(address))
            return false;

        PeerID id = this->addPeer(address);

        if (id < 0)
            return false;

        peers[id]->connected = false;

        this->send(id, EMT_ConnectRequest, 0, 0, true);

        return true;
    }


    //! The disconnect pendant
    //! Returns false if the peerID is invalid
    bool CUDPPort::disconnect(const PeerID id)
    {
        if (id < 0 || !peers[id])
            return false;

        //! we force a message here
        BitStream bstream;

        bstream.write<u8>(EMT_Disconnect);
        bstream.write<u16>(0);

        for (u16 u = 0; u < 20; u++)
        {
            if (this->sendRaw(peers[id]->address, bstream.getLength(), bstream.getData()))
                break;
        }

        if (!this->removePeer(id))
            return false;

        return true;
    }

    //! Disconnects all peers
    //! Returns number of peers disconnected
    u16 CUDPPort::disconnectAll(void)
    {
        u16 count = 0;

        for (u16 u = 0; u < maxPeers; u++)
        {
            count += this->disconnect(u);
        }

        return count;
    }

    bool CUDPPort::isConnected(const SNetworkAddress& addr)
    {
        return (this->getPeer(addr) != -1);
    }

    //! Gets the Id of a certain address
    PeerID CUDPPort::getPeer(const SNetworkAddress& addr)
    {
        for (u16 u = 0; u < maxPeers; u++)
        {
            if (peers[u] &&  peers[u]->address == addr )
                return u;
        }

        return -1;
    }

    //! Get the latency of a peer
    s32 CUDPPort::getLatency(const PeerID id)
    {
        if (!peers[id])
            return -1;

        return peers[id]->latency;
    }

    void CUDPPort::doPingPong(PeerID u)
    {
        if (peers[u]->pingSend && (getLocalTime() - peers[u]->pingTime) > static_cast<u32>(this->getLatency(u)*5))
        {
            //! Timeout
            this->getCallback()->onPeerDisconnect(u, ENI_ConnectionTimeout);

            peers.remove(u);
        }
        else if (!peers[u]->pingSend && (getLocalTime() - peers[u]->pingTime) > static_cast<u32>(this->getLatency(u)*2))
        {
            peers[u]->pingSend = true;
            peers[u]->pingTime = getLocalTime();

            //! Send the next ping
            this->send(u, EMT_Ping, 0, 0, true);
        }
    }

    //! Sends and receives data
    void CUDPPort::work(void)
    {
        //! Callback (most people wont use this)
        this->getCallback()->onUpdate();

        //! I diveded work() into these 4 functions
        //! Makes everything a lot more straightforward
        this->processReceiving();
        this->processSending();

        for (u16 u = 0; u < this->getMaxPeers(); u++)
        {
            if (peers[u])
            {
                this->doPingPong(u);
                this->doReliable(u);
            }
        }
    }

    void CUDPPort::doReliable(PeerID id)
    {
        //! How often should we do this?
        for (u16 u = 0; u < this->getMaxPeers(); u++)
        {
            if (peers[u])
            {
                std::vector<std::map<u16, u32>::iterator> removals;

                for (std::map<u16, u32>::iterator it = peers[u]->lastReceived.begin(); it != peers[u]->lastReceived.end(); it++)
                {
                    if (getLocalTime() - (*it).second > static_cast<u32>(this->getLatency(u))*10)
                        removals.push_back(it);
                }

                for (std::vector<std::map<u16, u32>::iterator>::iterator it = removals.begin(); it != removals.end(); it++)
                    peers[u]->lastReceived.erase(*it);
            }
        }
    }

    //! Get the protocol used
    EPortType CUDPPort::getType() const
    {
        return EPT_UserDatagramProtocol;
    }

    //! Send data to peer - Reliable or not
    bool CUDPPort::send(const PeerID peer, const u8 type, const u32 length, const c8* data, const bool reliable)
    {
        //! Message Protocol
        //! 1 Byte - Type (user defined)
        //! 2 Bytes - PacketID (if reliable)
        //! 1 Byte - Type (if reliable the first type will be used to say its reliable)
        //! 2 Bytes - Length
        //! Length Bytes - Data#

        //! -> length + 3/5 Bytes

        if (peer < 0 || !peers[peer])
            return false;

        //! If its reliable we do something different
        if (length > 1000)
        {
            u32 chunks = length / 1000;

            if (length%1000 > 0)
                chunks++;

            u16 pId = peers[peer]->lastPacketID+1;
            peers[peer]->lastPacketID = pId;

            BitStream bstream(data,length);

            //! This is the actual header
            //! Which has the id 0
            {
                BitStream outmsg;

                outmsg.write<u32>(0);
                outmsg.write<u32>(chunks);
                outmsg.write<u8>(type);

                this->internal_send(peer, EMT_BigMessage, outmsg.getLength(), outmsg.getData(), true, pId);
            }

            for (u32 u = 1; u < chunks; u++)
            {
                BitStream outmsg(1000+4);

                outmsg.write<u32>(u, false);
                bstream.copyRawArray<c8>(&outmsg, 1000, true, false);

                this->internal_send(peer, EMT_BigMessage, outmsg.getLength(), outmsg.getData(), true, pId);
            }

            //! The last message needs special treatment too
            {
                BitStream outmsg(length%1000+4);

                outmsg.write<u32>(chunks, false);
                bstream.copyRawArray<c8>(&outmsg, length%1000, true, false);

                this->internal_send(peer, EMT_BigMessage, outmsg.getLength(), outmsg.getData(), true, pId);
            }

            return true;
        }

        return this->internal_send(peer, type, length, data, reliable, -1);
    }

    //! Internal send fucntion you can set a reliable id here
    bool CUDPPort::internal_send(const PeerID peer, const u8 type, const u16 length, const c8* data, const bool reliable, s32 pId)
    {
        u16 headerSize = 3;

        if (reliable)
            headerSize += 3;

        BitStream *message = new BitStream(length+headerSize);

        if (reliable)
        {
            if (pId < 0)
            {
                pId = static_cast<u16>(peers[peer]->lastPacketID+1);
                peers[peer]->lastPacketID = pId;
            }

            //! We show the message is reliable and add its id
            message->write<u8>(EMT_Reliable, false);
            message->write<u16>(pId,false);
        }

        //! we add the type given by the user and the data
        message->write<u8>(type, false);
        message->write<u16>(length, false);

        //! You can also send empty messages (e.g. ping/pong)
        if (length > 0)
            message->writeRawArray<c8>(data, length, false);

        message->setByteIndex(0);

        if (reliable)
            messages.push_back(SUDPMessage(peer, pId, message));
        else
            messages.push_back(SUDPMessage(peer, -1, message));

        return true;
    }

    //! Send data to address
    bool CUDPPort::sendRaw(const SNetworkAddress& address, const u16 length, const c8* data)
    {
        //! Send function
        if (this->isIPv6())
        {
            sockaddr_in6 addr;

            if (getSockAddress6(addr, address) && sendto(socketfd, data, length, 0, (sockaddr*) &addr, sizeof(sockaddr_in6)))
                return true;
        }
        else
        {
            sockaddr_in addr;

            if (getSockAddress(addr, address) && sendto(socketfd, data, length, 0, (sockaddr*) &addr, sizeof(sockaddr_in)))
                return true;
        }

        if (errno != EWOULDBLOCK)
            perror("sendto");

        return false;
    }

    //! Query certain features
    bool CUDPPort::hasFeature(EFeatureType type) const
    {
        if (type == EFT_RawMessage)
            return true;
        else if (type == EFT_UnreliableMessage)
            return true;
        else if (type == EFT_ReliableMessage)
            return true;
        else if (type == EFT_Latency)
            return true;
        else if (type == EFT_BigMessage)
            return true;
        else
            return false;
    }

    u16 CUDPPort::getPeerCount()
    {
        u16 count = 0;

        for (u16 u = 0; u < maxPeers; u++)
        {
            if (peers[u])
                count++;
        }

        return count;
    }

    //! Remove a peer (!= disconnect)
    bool CUDPPort::removePeer(PeerID id)
    {
        if (id < 0 || id > maxPeers || !peers[id])
            return false;

        //! Search for messages we wont need anymore
        std::vector<std::vector<SUDPMessage>::iterator> removals;

        for (std::vector<SUDPMessage>::iterator it = messages.begin(); it != messages.end(); it++)
        {
            if ((*it).peer == id)
            {
                delete (*it).data;
                removals.push_back(it);
            }
        }

        for (std::vector<std::vector<SUDPMessage>::iterator>::iterator it = removals.begin(); it != removals.end(); it++)
            messages.erase(*it);

        delete peers[id];
        peers[id] = 0;

        return true;
    }

    //! Adds a peer (!= connect)
    PeerID CUDPPort::addPeer(SNetworkAddress address)
    {
        for (u16 u = 0; u < this->getMaxPeers(); u++)
        {
            if (!peers[u])
            {
                peers[u] = new SUDPPeer(address);

                return u;
            }
        }

        return -1;
    }

    //! Processes out messages
    void CUDPPort::processSending()
    {
        std::vector<std::vector<SUDPMessage>::iterator> removals;

        for (std::vector<SUDPMessage>::iterator it = messages.begin(); it != messages.end(); ++it)
        {
            if (!peers[(*it).peer] || !(*it).data)
            {
                if ((*it).data)
                    delete (*it).data;

                removals.push_back(it);
            }
            else
            {
                //! Resend if timeout
                if ((*it).id >= 0)
                {
                    if ((*it).timeSend < (getLocalTime() - static_cast<u32>(this->getLatency((*it).peer)*2))
                            && this->sendRaw(peers[(*it).peer]->address, (*it).data->getLength(), (*it).data->getData()))
                        (*it).timeSend = getLocalTime();
                }
                //! Send normal messages
                else if ((*it).id < 0)
                {
                    if (this->sendRaw(peers[(*it).peer]->address, (*it).data->getLength(), (*it).data->getData()))
                    {
                       if ((*it).data)
                            delete (*it).data;

                        removals.push_back(it);
                    }
                }
            }
        }

        for(std::vector< std::vector<SUDPMessage>::iterator >::iterator it = removals.begin(); it != removals.end(); ++it)
        {
            messages.erase((*it));
        }
    }

    //! Processes in messages
    void CUDPPort::processReceiving()
    {
        bool msgReceived = true;

        while (msgReceived)
        {
            //! First set it to false
            msgReceived = false;

            //! The address of the peer
            SNetworkAddress addr;

            //! The length read from the systems buffer
            s32 readLength = 0;

            //! Clear the buffer (global member)
            memset(&buffer[0], 0, 1024);

            //! The low level UDP code (either IPv4 or 6)
            if (this->isIPv6())
            {
                sockaddr_in6 cliAddr;
                socklen_t len = sizeof(cliAddr);

                readLength = recvfrom( socketfd, &buffer[0], 1024, 0, (struct sockaddr *) &cliAddr, &len );

                if (readLength > 0 && getNetworkAddress6(cliAddr, addr))
                    msgReceived = true;
            }
            else
            {
                sockaddr_in cliAddr;
                socklen_t len = sizeof(cliAddr);

                readLength = recvfrom( socketfd, &buffer[0], 1024, 0, (struct sockaddr *) &cliAddr, &len );

                if (readLength > 0 && getNetworkAddress(cliAddr, addr))
                    msgReceived = true;
            }

            if (msgReceived)
            {
                //!  Setup the BitStream
                BitStream data(&buffer[0], readLength);

                //! The id of the package (only if reliable)
                u16 id = 0;

                //! The real length (length + headerSize - see this->send())
                u16 realLength = 3;

                //! The message type (see EMessageType)
                EMessageType type ;

                //! The first byte either messageType or EMT_Reliable
                EMessageType first = static_cast<EMessageType>( data.read<u8>() );

                //! We get the peers id from the INetworkAddress
                PeerID peer = this->getPeer(addr);

                //! If reliable we have to read the id and the "real" type
                if (first == EMT_Reliable)
                {
                    id = data.read<u16>();
                    type = static_cast<EMessageType>( data.read<u8>() );

                    realLength += 3;
                }
                else
                    type = first;

                //!The length of the message
                u16 length = data.read<u16>();
                realLength += length;

                //! Check for corrupt package
                //! This has to be extended
                if (realLength == static_cast<u16>(readLength))
                {
                    if (type == EMT_ConnectRequest)
                    {
                        //! Maybe we send a request the same time?
                        //! And is this port actually listening for connections?
                        if (peer == -1 && this->isListening() && this->getPeerCount() < this->getMaxPeers() && callback->onConnectRequest(addr) )
                        {
                            PeerID u = this->addPeer(addr);

                            if (u >= 0)
                            {
                                peers[u]->connected = true;

                                callback->onPeerConnect(u);
                                this->send(u, EMT_ConnectAccepted, 0, 0, true);
                            }
                        }
                        else
                        {
                            //!Denie the connection
                            BitStream rawData;
                            rawData.write<c8>(EMT_ConnectDenied);
                            rawData.write<u16>(0);

                            this->sendRaw(addr, rawData.getLength(), rawData.getData());
                        }
                    }
                    else if (type == EMT_ArrivalNotification)
                    {
                        u16 notId = data.read<u16>();

                        //! Jippeee data was sent succesfully
                        //! Now we just have to delete the data from memory
                        if (notId >= 0 && peer >= 0 && peers[peer])
                        {
                            for (std::vector<SUDPMessage>::iterator it = messages.begin(); it !=messages.end(); it++)
                            {
                                if ((*it).peer == peer && (*it).id == notId)
                                {
                                    (*it).data->setByteIndex(0);

                                    //! If this is part of a big message only delete if were at the right pos
                                    if ((*it).data[0] == EMT_Reliable && (*it).data[1] == EMT_BigMessage)
                                    {
                                        (*it).data->setByteIndex(4);

                                        u32 this_pos = (*it).data->read<u32>();
                                        u32 real_pos = data.read<u32>();

                                        data.setByteIndex(data.getByteIndex() - 4);

                                        if (this_pos != real_pos)
                                            continue;
                                    }

                                    if ((*it).data)
                                        delete (*it).data;

                                    messages.erase(it);

                                    break;
                                }
                            }
                        }
                    }
                    else if (type == EMT_ConnectAccepted)
                    {
                        if (peer >= 0 && peers[peer] && !peers[peer]->connected)
                        {
                            peers[peer]->connected = true;
                            this->getCallback()->onPeerConnect(peer);
                        }
                    }
                    else if (type == EMT_ConnectDenied)
                    {
                        if (peer >= 0 && peers[peer] && !peers[peer]->connected)
                        {
                            callback->onPeerDisconnect(peer, ENI_ConnectionDenied);

                            //!delete afterwards so the user can acces the peer one last time
                            delete peers[peer];
                            peers[peer] = 0;
                        }

                    }
                    else if ( type == EMT_Ping)
                    {
                        this->send(peer, EMT_Pong, 0, 0, true);
                    }
                    else if ( type == EMT_Pong)
                    {
                        if (peer >= 0 &&peers[peer])
                        {
                            peers[peer]->latency = getLocalTime() - peers[peer]->pingSend;
                            peers[peer]->pingSend = false;
                            peers[peer]->pingTime = getLocalTime();
                        }
                    }
                    else if (type == EMT_Disconnect)
                    {
                        if (peer >= 0 && peers[peer])
                        {
                            callback->onPeerDisconnect(peer, ENI_ConnectionClosed);

                            //!delete afterwards so the user can acces the peer one last time
                            this->removePeer(peer);
                        }
                    }
                    else if (peer >= 0 &&peers[peer])
                    {
                        //! Actual data for the Callback! wohoooo

                        if (first == EMT_Reliable && type == EMT_BigMessage)
                        {
                            std::vector<SUDPBigMessage>::iterator it = peers[peer]->bigMessages.begin();

                            for (; it != peers[peer]->bigMessages.end(); ++it)
                            {
                                if ((*it).reliableId == id)
                                    break;
                            }

                            if (it == peers[peer]->bigMessages.end())
                                it = peers[peer]->bigMessages.insert(it, SUDPBigMessage(id));

                            u32 pos = data.read<u32>();

                            //! Notify that the chunk arrived
                            BitStream outMsg;
                            outMsg.write<u16>(id);
                            outMsg.write<u32>(pos);

                            this->send(peer, EMT_ArrivalNotification, outMsg.getLength(), outMsg.getData(), false);

                            if (pos == 0)
                            {
                                //! The first chunk tells us about the size etc...
                                (*it).size = data.read<u32>();
                                (*it).messageType = data.read<u8>();
                            }
                            else
                            {
                                u32 length = data.getLength() - data.getByteIndex();

                                BitStream chunk;
                                data.copyRawArray<c8>(&chunk, length, true, true);

                                (*it).chunks[pos] = chunk;

                                //! Has everything arrived?
                                if ((*it).chunks.size() == (*it).size)
                                {
                                    BitStream fullMsg;

                                    for (u32 u = 1; u <= (*it).size; u++)
                                    {
                                        fullMsg += (*it).chunks[u];
                                    }

                                    this->getCallback()->onMessageReceived(peer, (*it).messageType, fullMsg.getLength(), fullMsg.getData(), true);

                                    peers[peer]->bigMessages.erase(it);
                                }
                            }
                        }
                        else
                        {
                            c8* message = data.readRawArray<c8>(length);

                            if (first == EMT_Reliable)
                            {
                                //! It could happen that the package was already received so we just flush it down the toilet
                                bool alreadyReceived = false;

                                for (std::map<u16, u32>::iterator it = peers[peer]->lastReceived.begin(); it != peers[peer]->lastReceived.end() && !alreadyReceived; it++)
                                {
                                    if ((*it).first == id)
                                        alreadyReceived = true;
                                }

                                //! If its the first time we send a notification and the callback stuff
                                if (!alreadyReceived)
                                {
                                    BitStream outMsg;
                                    outMsg.write<u16>(id);

                                    this->send(peer, EMT_ArrivalNotification, outMsg.getLength(), outMsg.getData(), false);
                                    this->getCallback()->onMessageReceived(peer, type, length, message, true);
                                }
                            }
                            else
                                this->getCallback()->onMessageReceived(peer, type, length, message, false);

                            //! Now free memory
                            delete []message;
                        }
                    }
                }
            }
        }
    }

} //! end namespace onet
